import pandas as pd

PERIOD_FIVE_DAY = 5
PERIOD_TEN_DAY = 10
RSI_PERIOD = 14
M_PERIOD = 20 # 布林带以20日为指标
N_PARAMS = 2

def RSI(close: pd.DataFrame, periods=RSI_PERIOD):# 计算价格变动
    delta = close.diff()
    delta = delta.dropna()

    # 分类上涨和下跌
    up = delta.copy()
    down = delta.copy()
    up[up < 0] = 0
    down[down > 0] = 0

    # 计算平均上涨收益和平均下跌收益
    avg_up = up.rolling(window=periods).mean()
    avg_down = down.rolling(window=periods).mean().abs()

    # 计算RSI指标
    rs = avg_up / avg_down
    rsi = 100 - (100 / (1 + rs))
    return rsi

def add_factor(dir: str):
    df: pd.DataFrame = pd.read_csv(dir)
    assert(isinstance(df, pd.DataFrame))
    # get moving average, in 5 days period or 10 days period
    df['MA_5'] = df['close'].copy().rolling(window=PERIOD_FIVE_DAY).mean()
    df['MA_5'].fillna(value=0, inplace=True)
    df['MA_10'] = df['close'].rolling(window=PERIOD_TEN_DAY).mean()
    df['MA_10'].fillna(value=0, inplace=True)

    # to get rsi
    df['RSI'] = RSI(df['close'].copy(deep=True))
    df['RSI'].fillna(value=0, inplace=True)

    # to get bollin
    df['MEDIUM'] = df['close'].copy().rolling(window=M_PERIOD).mean()
    df['MEDIUM'].fillna(value=0, inplace=True)
 
    df['STD'] = df['close'].copy().rolling(M_PERIOD, min_periods=1).std(ddof=0)  # ddof代表标准差自由度
    df['STD'].fillna(value=0, inplace=True)
    df['UPPER'] = df['MEDIUM'] + N_PARAMS * df['STD']
    df['LOWER'] = df['MEDIUM'] - N_PARAMS * df['STD']

    diff = df['close'].copy() - df['close'].copy().shift(1)
    diff.fillna(value=0, inplace=True)
    diff[diff > 0] = 1
    diff[diff < 0] = -1
    df['DIRECTION'] = diff

    df.to_csv(dir, sep=',', index=False, header=True)
